import json
import os
import glob

def extract_content_from_files(input_dir, output_file):
    """
    从多个JSONL或JSON文件中提取所有"content"或"value"字段，并生成新的JSONL文件
    
    参数:
        input_dir: 包含输入文件的目录
        output_file: 输出的JSONL文件路径
    """
    # 获取输入目录中的所有JSONL和JSON文件
    jsonl_files = glob.glob(os.path.join(input_dir, "*.jsonl"))
    json_files = glob.glob(os.path.join(input_dir, "*.json"))
    
    # 检查是否有文件
    if not jsonl_files and not json_files:
        print(f"在目录 {input_dir} 中没有找到JSONL或JSON文件")
        return
    
    # 打开输出文件
    with open(output_file, 'w', encoding='utf-8') as outfile:
        # 处理JSONL文件
        for file_path in jsonl_files:
            print(f"处理JSONL文件: {file_path}")
            process_jsonl_file(file_path, outfile)
        
        # 处理JSON文件
        for file_path in json_files:
            print(f"处理JSON文件: {file_path}")
            process_json_file(file_path, outfile)
    
    print(f"处理完成，结果已保存到: {output_file}")

def process_jsonl_file(file_path, outfile):
    """处理单个JSONL文件"""
    with open(file_path, 'r', encoding='utf-8') as infile:
        for line in infile:
            try:
                data = json.loads(line.strip())
                extract_and_write_content(data, outfile)
            except json.JSONDecodeError as e:
                print(f"解析JSONL文件 {file_path} 时出错: {e}")
                continue

def process_json_file(file_path, outfile):
    """处理单个JSON文件"""
    with open(file_path, 'r', encoding='utf-8') as infile:
        try:
            data = json.load(infile)
            
            # 检查data是否是列表
            if isinstance(data, list):
                for item in data:
                    extract_and_write_content(item, outfile)
            else:
                extract_and_write_content(data, outfile)
                
        except json.JSONDecodeError as e:
            print(f"解析JSON文件 {file_path} 时出错: {e}")

def extract_and_write_content(data, outfile):
    """从数据中提取content或value字段并写入输出文件"""
    # 检查data是否包含conversations
    conversations = data.get("conversations", [])
    
    for conversation in conversations:
        # 尝试提取content字段
        content = conversation.get("content")
        if content:
            output_data = {"text": f"<s>{content}</s>"}
            json.dump(output_data, outfile)
            outfile.write('\n')
            continue
        
        # 尝试提取value字段
        value = conversation.get("value")
        if value:
            output_data = {"text": f"<s>{value}</s>"}
            json.dump(output_data, outfile)
            outfile.write('\n')

if __name__ == "__main__":
    # 输入目录（包含JSONL和JSON文件的目录）
    input_directory = "dataset/Cinego"
    
    # 输出文件路径
    output_jsonl_path = "tokenizer_video_data.jsonl"
    
    # 执行提取
    extract_content_from_files(input_directory, output_jsonl_path)